/*-------------------------------------------------------------------------
 *
 * txn_timestmap.c
 *  timestamp coordination of transactions
 * 
 * Copyright (c) 2020, Alibaba Group Holding Limited
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. *
 * src/backend/distributed_txn/txn_timestamp.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/xact.h"
#include "access/transam.h"
#include "access/twophase.h"
#include "distributed_txn/txn_timestamp.h"
#include "fmgr.h"
#include "libpq/libpq.h"
#include "utils/builtins.h"
#include "access/mvccvars.h"
#include "storage/proc.h"
#include "storage/procarray.h"

/* user-set guc parameter */
bool enable_timestamp_debug_print = false;
/* 
 * start timestamp used in both coordinator and worker:
 */
static bool backendReceivedTimestamp = false;
/* start_ts is generated by node itself or coordinated */
static LogicalTime txnStartTs = 0;
/* commit_ts = ClockTick() */
static LogicalTime txnCommitTs = 0;
/* coordinated_ts = Max(prepare_ts) */
static LogicalTime txnCoordinatedCommitTs = 0;
static LogicalTime replyTs = 0;

IsCoordinating2PCHook IsCoordinating2PC = NULL;

static void TxnSetStartTs(LogicalTime);
static void TxnSetCommitTs(LogicalTime);

LogicalTime TxnGetStartTs(void)
{
    return txnStartTs;
}

LogicalTime TxnGetCoordinatedCommitTs(void)
{
    if (enable_timestamp_debug_print)
        elog(LOG, "Get coordinated committs "UINT64_FORMAT" to send to workers", txnCoordinatedCommitTs);

    return txnCoordinatedCommitTs;
}

/*
 * Called by GetSnapshotData on Coordinator or Worker
 * On coordinator: 
 *      use ClockNow(), since GetSnapshotData is determined by XactIsoLevel
 * On worker: 
 *      in RC/RR,  receiving start_ts happens before GetSnapshotData
 * We rule them all in following ways:
 * 1. TxnGenerateStartTs() initializes txnStartTs only at the transaction start
 * 2. BackendRecvTimestamp() updates txnStartTs when receving timestamp from coordinator,
 *      so the next time GetSnapshotData() will return the latest txnStartTs
 */
LogicalTime TxnGetOrGenerateStartTs(bool latest)
{
     LogicalTime res;

    if (!txnUseGlobalSnapshot || latest)
    {
        res = LogicalClockNow();
        if (enable_timestamp_debug_print)
        {
            elog(LOG, "PhysicalNow: "LOGICALTIME_FORMAT, 
                        LOGICALTIME_STRING(PhysicalClockNow()));
            const char *snapshotFreshness = latest ? "latest" : "stable";
            const char *snapshotCreator = txnUseGlobalSnapshot ? "global" : "local";
            elog(LOG, "Generate start_ts "UINT64_FORMAT LOGICALTIME_FORMAT 
                        " use %s %s snapshot MyTmin "UINT64_FORMAT " procno %d",
                        res, LOGICALTIME_STRING(res), snapshotFreshness, snapshotCreator,
                        pg_atomic_read_u64(&MyPgXact->tmin), MyProc->pgprocno);
        }
        if (!latest)
            txnStartTs = res;
    }
    else
    {
        res = txnStartTs;
        if (enable_timestamp_debug_print)
        {
            const char *snapshotCreator = txnUseGlobalSnapshot ? "global" : "local";
            elog(LOG, "Get start_ts "UINT64_FORMAT LOGICALTIME_FORMAT " from coordinator "
                       "using %s snapshot MyTmin "UINT64_FORMAT " procno %d", 
                       res, LOGICALTIME_STRING(res), snapshotCreator, pg_atomic_read_u64(&MyPgXact->tmin), MyProc->pgprocno);
        }
    }
    return res;
}

/**
 * Called by TransactionIdAsyncCommitTree() on Coordinator or Worker
 * So it needs to meet the requirements of 2PC, 1PC, 0PC on both coordinator and worker.
 * 
 * On coordinator: 
 *      use coordinated_commit_ts if in 2PC transaction, else use ClockTick()
 * On worker: 
 *      use commit_ts assigned from txn_commit_prepare() if in 2PC transaction,
 *      else use ClockTick()
 */
LogicalTime TxnGetOrGenerateCommitTs(bool fromCoordinator)
{
    LogicalTime res;
    
    if (txnCoordinatedCommitTs != 0)
    {
        if (txnUseGlobalSnapshot && false == fromCoordinator)
            elog(ERROR, "Coordinator did not pass any commit timestamp down");
        
        /* coordinated ts from 2PC */
        res = txnCoordinatedCommitTs;
        if (enable_timestamp_debug_print)
            elog(LOG, "Get commit timestamp from coordinator "UINT64_FORMAT " logical clock "UINT64_FORMAT, 
                                res, ToLogicalClock(res));
    }
    else 
    {
        if (txnUseGlobalSnapshot && fromCoordinator)
            elog(ERROR, "Coordinator should not pass commit timestamp down for 1PC transactions");
        
        /* generate commit_ts by myself */
        if (txnCommitTs)
            elog(ERROR, "commit timestamp exists unexpectly "UINT64_FORMAT, txnCommitTs);

        TxnSetCommitTs(LogicalClockTick());
        res = txnCommitTs;
        if (enable_timestamp_debug_print)
            elog(LOG, "Get commit timestamp locally "UINT64_FORMAT " logical clock "UINT64_FORMAT, 
                                res, ToLogicalClock(res));
    }
    Assert(!COMMITSEQNO_IS_SUBTRANS(res));
    return res;
}

void AtEOXact_txn(void) 
{
    if (enable_timestamp_debug_print)
        elog(LOG, "AtEOXact_txn");

    txnUseGlobalSnapshot = false;
    RecentGlobalTs = InvalidCommitSeqNo;
    txnCoordinatedCommitTs = 0;
    TxnSetStartTs(0);
    TxnSetCommitTs(0);
}

static void TxnSetStartTs(LogicalTime startTs)
{
    txnStartTs = startTs;
    if (startTs)
        ereport(DEBUG1, (errmsg("set start_ts %lu", startTs)));
}

static void TxnSetCommitTs(LogicalTime commitTs)
{
    txnCommitTs = commitTs;
  
    if (enable_timestamp_debug_print)
        elog(LOG, "set commit timestamp "UINT64_FORMAT, commitTs);

}

void TxnSetCoordinatedCommitTs(LogicalTime ts)
{
    Assert(txnCoordinatedCommitTs == 0 || ts == 0);
    LogicalClockUpdate(ts);
    txnCoordinatedCommitTs = ts;
    
    if (enable_timestamp_debug_print)
        elog(LOG, "set commit timestamp from coordinator "UINT64_FORMAT, ts);

}

void TxnSetCoordinatedCommitTsFromStr(const char *ts)
{
    LogicalTime commit_ts = (LogicalTime)strtoull(ts, NULL, 0);
    if (commit_ts == 0) 
        ereport(ERROR, (errmsg("coordinated commit_ts is 0")));
    
    TxnSetCoordinatedCommitTs(commit_ts);
}

void TxnSetReplyTimestamp(LogicalTime ts)
{
    replyTs = ts;
}

LogicalTime TxnGetAndClearReplyTimestamp(void)
{
    LogicalTime res = replyTs;

    if (!backendReceivedTimestamp)
    {
        return 0;
    }
    replyTs = 0;
   
    if (enable_timestamp_debug_print)
        elog(LOG, "reply timestamp "UINT64_FORMAT " logical clock "UINT64_FORMAT, 
                                res, ToLogicalClock(res));

    return res;
}

/*
 * Receive timestamp from coordinator
 * read-committed: receive start_ts every statement, use start_ts as snapshot.csn
 * repeatable-read: receive start_ts before first statement
 * 0PC: no explicit transaction exists, just update the clock but not set snapshot
*/
void BackendRecvTimestamp(LogicalTime ts)
{
    /*
	 * Be carefull of the order between setting MyPgXact->tmin
	 * and acquiring the ts_lock to fetch maxCommitTs
	 * which is critical to the correctness of garbage collection algorithm.
	 * Written by Junbin Kang, 2020.01.20
	 */
	pg_atomic_write_u64(&MyPgXact->tmin, ts);
	pg_memory_barrier();

    backendReceivedTimestamp = true;
    RecentGlobalTs = LogicalClockUpdate(ts);
    TxnSetStartTs(ts);
    txnUseGlobalSnapshot = true;    
    if (enable_timestamp_debug_print)
        elog(LOG, "recv start timestamp trans %d "UINT64_FORMAT " procno %d", IsTransactionState(), ts, MyProc->pgprocno);
}

/*
 * This timestamp from worker could be either prepare_ts or commit_ts
 * prepare_ts: coordinator in COORD_TRANS_STARTED state and sended prepare command
 * commit_ts: any other scenarios
*/
void FrontendRecvTimestamp(LogicalTime ts)
{
    LogicalClockUpdate(ts);
    if (enable_timestamp_debug_print)
        elog(LOG, "recv reply timestamp trans "UINT64_FORMAT " logical clock "UINT64_FORMAT, 
                    ts, ToLogicalClock(ts));
}

PG_FUNCTION_INFO_V1(txn_get_start_ts);
PG_FUNCTION_INFO_V1(txn_commit_prepared);

/**
 * inspect start_ts of a transaction, for test purpose
 */
Datum
txn_get_start_ts(PG_FUNCTION_ARGS)
{
    LogicalTime ts = TxnGetStartTs();
    PG_RETURN_UINT64(ts);    
}

/**
 * Commit a prepared transaction, with timestamp assigned from coordinator
 */
Datum 
txn_commit_prepared(PG_FUNCTION_ARGS)
{
	const char *gid = text_to_cstring(PG_GETARG_TEXT_PP(0));
    LogicalTime commit_ts = PG_GETARG_INT64(1);
    
    if (commit_ts == 0) 
        ereport(ERROR, (errmsg("coordinated commit_ts is 0")));
    
    TxnSetCoordinatedCommitTs(commit_ts);
    FinishPreparedTransaction(gid, true);
    PG_RETURN_NULL();
}
